

#include "main.h"

/******************************************************************************/

// 一条stream以accept获得的port作为唯一标记
// 在listen端, fd是accept获得的socket
// 在forward端, fd是用于connect的socket
// 一条stream要在一个tunnel上传输

typedef struct stream_t
{
	int fd;
	int port;
	void *tunnel;
}STREAM;

// stream的数据头部
typedef struct pkt_header_t {
	u16 port;
	u8  code;
	u8  unuse;
}PKT_HDR;


#define MAX_STREAMS  64

STREAM streams[MAX_STREAMS];

SOCKET listen_fd = 0;


/******************************************************************************/


STREAM *new_stream(int fd, int port, void *tunnel)
{
	int i;

	for(i=0; i<MAX_STREAMS; i++){
		if(streams[i].fd==0){
			streams[i].fd = fd;
			streams[i].port = port;
			streams[i].tunnel = tunnel;
			return &streams[i];
		}
	}

	printf("Too many streams!\n");
	return NULL;
}


int remove_stream(STREAM *st)
{
	shutdown(st->fd, SD_BOTH);
	closesocket(st->fd);

	evio_del(st->fd);

	st->fd = 0;
	st->port = 0;

	return 0;
}


STREAM *find_stream(int port)
{
	int i;

	for(i=0; i<MAX_STREAMS; i++){
		if(streams[i].fd && streams[i].port==port){
			return &streams[i];
		}
	}

	return NULL;
}


void close_tunnel_stream(void *tunnel)
{
	int i;

	for(i=0; i<MAX_STREAMS; i++){
		if(streams[i].fd && streams[i].tunnel == tunnel){
			remove_stream(&streams[i]);
		}
	}
}


int socket_nonblock(SOCKET skt)
{
#if 1

#ifdef __WIN32__
	u_long iMode = 1;
	return ioctlsocket(skt, FIONBIO, &iMode);
#else
	int flags = fcntl(skt, F_GETFL, 0);
	if(flags<0)
		flags = 0;
	fcntl(skt, F_SETFL, flags|O_NONBLOCK);
	return 0;
#endif

#else
	return 0;
#endif
}


/******************************************************************************/


int stream_send(STREAM *st, int code, const void *buf, int len)
{
	char hbuf[16];
	char *sbuf = (char*)buf;
	PKT_HDR *ph = (PKT_HDR*)sbuf;

	if(sbuf)
		sbuf -= sizeof(PKT_HDR);
	else
		sbuf = hbuf;
	ph = (PKT_HDR*)sbuf;

	memset(ph, 0, sizeof(PKT_HDR));
	ph->port = st->port;
	ph->code = code;

	return kcp_send(st->tunnel, sbuf, sizeof(PKT_HDR)+len);
}


int handle_stream(STREAM *st)
{
	char pbuf[1536];
	char *rbuf = pbuf+32;

	// 有数据到来
	int len = recv(st->fd, rbuf, gcfg.mtu, 0);
	if(len<=0){
		if(len<0){
			printf("stream[%d]: recv failed! %d\n", st->port, WSAGetLastError());
		}else if(len==0){
			printf("stream[%d]: connection closed!\n", st->port);
		}
		stream_send(st, PCODE_DISCONNECT, NULL, 0);
		remove_stream(st);
		return 0;
	}
	//printf("[%06"LLD"]stream[%d]: new packet len=%d\n", get_ticks(), st->port, len);
	//hex_dump(NULL, rbuf, 64);

	stream_send(st, PCODE_DATA, rbuf, len);
	return 0;
}


int handle_connect(int src_port, void *tunnel)
{
	struct sockaddr_in bind_addr;
	int nfd, retv, connect_port;

	nfd = socket(AF_INET, SOCK_STREAM, 0);
	if (nfd==INVALID_SOCKET) {
		printf("target socket create error!\n");
		return -1;
	}

	retv = 2*1024*1024;
	setsockopt (nfd, SOL_SOCKET, SO_SNDBUF, (char*)&retv, sizeof(retv));
	setsockopt (nfd, SOL_SOCKET, SO_RCVBUF, (char*)&retv, sizeof(retv));

	bind_addr.sin_family = AF_INET;
	bind_addr.sin_port = htons(gcfg.forward_port);
	if(gcfg.forward_ip[0])
		bind_addr.sin_addr.s_addr = inet_addr(gcfg.forward_ip);
	else
		bind_addr.sin_addr.s_addr = INADDR_ANY;

	printf("\nConnect: client:%d --> %s:%d ...\n", src_port, gcfg.forward_ip, gcfg.forward_port);
	retv = connect(nfd, (struct sockaddr *)&bind_addr, sizeof(bind_addr));
	if(retv){
		printf("connect failed! %d\n", WSAGetLastError());
		return -1;
	}
	socket_nonblock(nfd);

	socklen_t alen = sizeof(bind_addr);
	getsockname(nfd, (struct sockaddr *)&bind_addr, &alen);
	connect_port = ntohs(bind_addr.sin_port);
	printf("    use socket %d\n\n", connect_port);

	STREAM *st = new_stream(nfd, src_port, tunnel);
	evio_add(nfd, handle_stream, NULL, NULL, st);
	evio_enable(nfd, EV_IN);
	return 0;
}


void stream_recv(void *tunnel, const void *buf, int len)
{
	int retv;
	STREAM *st;
	PKT_HDR *ph = (PKT_HDR*)buf;

	buf += sizeof(PKT_HDR);
	len -= sizeof(PKT_HDR);

	//printf("[%06"LLD"]stream_recv: code=%02x port=%d len=%d\n", get_ticks(), ph->code, ph->port, len);


	if(ph->code==PCODE_MSG){
		printf("\nPeer MSG: %s\n", (char*)buf);
	}else if(ph->code==PCODE_CONNECT){
		handle_connect(ph->port, tunnel);
	}else if(ph->code==PCODE_DISCONNECT){
		st = find_stream(ph->port);
		if(st){
			printf("Disconnect: port=%d\n", st->port);
			remove_stream(st);
		}
	}else if(ph->code==PCODE_DATA){
		st = find_stream(ph->port);
		const char *sp = (const char*)buf;
		while(st && len>0){
			retv = send(st->fd, sp, len, 0);
			if(retv>=0){
				len -= retv;
				sp += retv;
			}else{
				printf("stream_recv[%d]: send to socket failed! %d\n", st->port, WSAGetLastError());
				break;
			}
		}
	}

}


/******************************************************************************/


int handle_accept(void)
{
	struct sockaddr_in bind_addr;
	STREAM *st;

	// 有新的连接
	socklen_t alen = sizeof(bind_addr);
	int nfd = accept(listen_fd, (struct sockaddr *)&bind_addr, &alen);
	if(nfd<0){
		printf("accept failed! %d\n", WSAGetLastError());
		return -1;
	}
	socket_nonblock(nfd);

	int retv = 2*1024*1024;
	setsockopt (nfd, SOL_SOCKET, SO_SNDBUF, (char*)&retv, sizeof(retv));
	setsockopt (nfd, SOL_SOCKET, SO_RCVBUF, (char*)&retv, sizeof(retv));

	int nport = ntohs(bind_addr.sin_port);
	printf("\nPort[%d]: New connection from %s:%d\n\n", gcfg.listen_port, inet_ntoa(bind_addr.sin_addr), nport);

	// 加入监听队列
	st = new_stream(nfd, nport, rawio_listen_tun());
	evio_add(nfd, handle_stream, NULL, NULL, st);
	evio_enable(nfd, EV_IN);

	// 通知server建立target连接
	stream_send(st, PCODE_CONNECT, NULL, 0);

	return 0;
}


int create_listen_port(void)
{
	int retv;
	struct sockaddr_in bind_addr;

	// 创建监听socket
	listen_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (listen_fd==INVALID_SOCKET) {
		printf("listen socket create error!\n");
		return -1;
	}

	char on = 1;
	setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
	//socket_nonblock(listen_fd);

	bind_addr.sin_family = AF_INET;
	bind_addr.sin_port = htons(gcfg.listen_port);
	if(gcfg.listen_ip[0])
		bind_addr.sin_addr.s_addr = inet_addr(gcfg.listen_ip);
	else
		bind_addr.sin_addr.s_addr = INADDR_ANY;

	retv = bind(listen_fd, (struct sockaddr *)&bind_addr, sizeof(bind_addr));
	if(retv){
		printf("bind failed! %d\n", WSAGetLastError());
		return -1;
	}

	printf("kcp2raw listen on %s:%d ...\n", gcfg.listen_ip, gcfg.listen_port);
	retv = listen(listen_fd, 64);
	if(retv){
		printf("listen failed! %d\n", WSAGetLastError());
		return -1;
	}

	evio_add(listen_fd, handle_accept, NULL, NULL, NULL);
	evio_enable(listen_fd, EV_IN);
	return 0;
}


void handle_tmout(int ms)
{
	int retv;

	// 反向代理时,只有客户端连接后才开始listen
	if(gcfg.listen_port && listen_fd==0 && rawio_listen_tun()){
		retv = create_listen_port();
		if(retv)
			evio_exit();
	}

	rawio_heartbeat(ms);
}


int run_tcp_server(void)
{
	int retv;

	evio_init(100);

	memset(streams, 0, sizeof(streams));
	retv = rawio_init(stream_recv, NULL);
	if(retv)
		return retv;

	evio_loop(handle_tmout);

	return 0;
}


/******************************************************************************/

